from functools import wraps

def support_predict(func):
    """让feature方法支持 is_predict 参数装饰器"""
    @wraps(func)
    def wrapper(df, is_predict=False):
        # 训练模式：保持原有逻辑不变
        if not is_predict:
            return func(df)
        
        # 预测模式：只返回特征矩阵 X
        result = func(df)
        if isinstance(result, tuple):
            return result[0]
        return result
    return wrapper 